python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py (2)

151-224: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

No dedicated unit tests for the new blockscaled descriptor-packing logic.

make_blockscaled_instr_desc/blockscaled_mma_op_to_idesc implement intricate bit-packing for the SM100 MXF8F6F4 idesc, but no test file analogous to test_DSA_mxfp8_scale_utils.py exercises this packing directly (only indirectly through downstream kernel correctness tests, if any). A focused unit test asserting expected bit values for representative (M, N, sf_id, major) combinations would catch regressions independent of full kernel execution.

As per path instructions for python/cudnn/**, which say to check "whether there are test cases in test/python/fe_api".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py` around lines
151 - 224, Add focused unit coverage for the new blockscaled descriptor packing
by creating tests under test/python/fe_api for make_blockscaled_instr_desc and
blockscaled_mma_op_to_idesc. Exercise representative M/N, sf_id, and major
combinations, and assert the returned descriptor has the expected bit fields, so
the SM100 MXF8F6F4 packing logic is validated directly instead of only through
downstream kernel tests.

---

82-89: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Fallback getattr(..., None) could silently misclassify if cutlass_type is unexpectedly None.

If none of the type attributes exist in the installed cutlass module and cutlass_type itself happens to be None, the first is None comparison would incorrectly match and return MXF8F6F4Format.E4M3 instead of raising the intended TypeError. This is unlikely in normal use (real dtype objects are always passed) but the redundant dual-name checks (FloatE4M3FN/Float8E4M3FN) suggest uncertainty about which name is available in the installed cutlass version.

🛡️ Safer alternative

-    if cutlass_type is getattr(cutlass, "FloatE4M3FN", None):
-        return MXF8F6F4Format.E4M3
-    if cutlass_type is getattr(cutlass, "Float8E4M3FN", None):
-        return MXF8F6F4Format.E4M3
-    if cutlass_type is getattr(cutlass, "FloatE5M2", None):
-        return MXF8F6F4Format.E5M2
-    if cutlass_type is getattr(cutlass, "Float8E5M2", None):
-        return MXF8F6F4Format.E5M2
+    e4m3_names = ("FloatE4M3FN", "Float8E4M3FN")
+    e5m2_names = ("FloatE5M2", "Float8E5M2")
+    if any(cutlass_type is getattr(cutlass, name, object()) for name in e4m3_names):
+        return MXF8F6F4Format.E4M3
+    if any(cutlass_type is getattr(cutlass, name, object()) for name in e5m2_names):
+        return MXF8F6F4Format.E5M2

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py` around lines
82 - 89, The dtype-to-format mapping in the helper that checks `cutlass_type`
against `FloatE4M3FN`, `Float8E4M3FN`, `FloatE5M2`, and `Float8E5M2` can
misclassify a `None` input because `getattr(..., None)` makes the first
comparison match accidentally. Update this conversion logic to avoid treating a
missing cutlass symbol or a `None` dtype as a valid match, and keep the intended
`TypeError` path in place for unsupported inputs while still handling the
alternate cutlass type names.

python/cudnn/deepseek_sparse_attention/utils/tensor_conversion.py (1)

15-32: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Docstring doesn't cover the new parameter annotations.

to_cute_tensor now has full type hints but the docstring still only describes the return in one line; assumed_align, leading_dim, fully_dynamic, enable_tvm_ffi, and divisibility aren't documented.

As per path instructions for python/cudnn/**, which say to "Focus on documentation."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudnn/deepseek_sparse_attention/utils/tensor_conversion.py` around
lines 15 - 32, The to_cute_tensor docstring is missing documentation for its
newly annotated parameters. Update the docstring in to_cute_tensor to describe
assumed_align, leading_dim, fully_dynamic, enable_tvm_ffi, and divisibility,
while keeping the return description accurate; use the existing function
signature in tensor_conversion.py as the reference point.

python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py (1)

52-98: 📐 Maintainability & Code Quality | 🔵 Trivial

Missing test coverage for the generic FP8 epilogue and varlen+FP8 on SM90.

The new test_DSA_indexer_forward_wrapper_fp8_sm90_matches_dequant_reference test only exercises qhead_per_kv_head=64 with divisible shapes, which routes through use_fast_qh64_epilogue/use_unchecked_qh64_full/use_unchecked_qh64_masked (lines 95-97). The generic _epilogue_store_to_gmem FP8 path (qhead_per_kv_head=32, or shapes that don't satisfy the "unchecked" divisibility conditions) and the varlen (THD) + FP8 combination on SM90 appear untested in test/python/fe_api/dsa/.

Given the amount of new FP8-specific branching (split-warpgroup producer/consumer, rKScale shared-memory scale application, LSE accumulation), consider adding at least one test with qhead_per_kv_head=32 and one varlen+FP8 SM90 test.

Based on path instructions: "Focus on whether there are test cases in test/python/fe_api" for python/cudnn/** files.

Also applies to: 369-535

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py`
around lines 52 - 98, Add test coverage in test/python/fe_api for the SM90 FP8
forward path beyond the current qhead_per_kvhead=64 fast epilogue case: add one
case that forces the generic _epilogue_store_to_gmem branch by using
qhead_per_kvhead=32 or a shape that does not satisfy
use_unchecked_qh64_full/use_unchecked_qh64_masked in IndexerFwdSM90, and add one
varlen (THD) + FP8 SM90 test that exercises the is_varlen path with FP8
scales/LSE handling. Keep the existing reference-matching style used by
test_DSA_indexer_forward_wrapper_fp8_sm90_matches_dequant_reference so the new
cases verify the generic FP8 epilogue and varlen+FP8 behavior.

Source: Path instructions

python/cudnn/deepseek_sparse_attention/indexer_forward/api.py (1)

209-215: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the newly added public parameters. The indexer_forward_wrapper docstring still only describes {'scores'} and q_causal_offsets, but the signature now exposes precision, q_scale, k_scale, sf_vec_size, return_lse, and lse_out. It's also worth noting that return_lse/lse_out are only honored on SM90 (SM100 raises NotImplementedError) so callers aren't surprised. As per path instructions ("Focus on documentation").

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cudnn/deepseek_sparse_attention/indexer_forward/api.py` around lines
209 - 215, Update the indexer_forward_wrapper docstring to document all newly
exposed public parameters: precision, q_scale, k_scale, sf_vec_size, return_lse,
and lse_out, along with any expected behavior for the outputs. Also note in the
wrapper docs that return_lse and lse_out are only supported on SM90, while SM100
raises NotImplementedError, so callers understand the hardware-specific
limitation.

Source: Path instructions